package com.example.nioprc.infrastructure.server;

import com.example.nioprc.infrastructure.protocol.RpcRequest;
import com.example.nioprc.infrastructure.protocol.RpcResponse;
import com.example.nioprc.infrastructure.registry.RpcServiceRegistry;
import com.example.nioprc.infrastructure.registry.ServiceRegistry;
import com.example.nioprc.infrastructure.serializer.JsonSerializer;
import com.example.nioprc.infrastructure.serializer.Serializer;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.lang.reflect.InvocationTargetException;
import java.lang.reflect.Method;
import java.net.InetSocketAddress;
import java.nio.ByteBuffer;
import java.nio.channels.*;
import java.util.Iterator;
import java.util.Map;
import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicLong;

/**
 * 基于Java原生NIO的RPC服务器实现
 */
public class NioRpcServer implements RpcServer {
    private static final Logger logger = LoggerFactory.getLogger(NioRpcServer.class);

    private final int port;
    private final ServiceRegistry serviceRegistry;
    private final Serializer serializer;
    private final ExecutorService threadPool;
    private final Map<SocketChannel, ConnectionState> connectionStates;
    private final AtomicLong requestCounter = new AtomicLong(0);

    private ServerSocketChannel serverSocketChannel;
    private Selector selector;
    private volatile boolean running = false;

    public NioRpcServer(int port) {
        this(port, new RpcServiceRegistry(), new JsonSerializer());
    }

    public NioRpcServer(int port, ServiceRegistry serviceRegistry, Serializer serializer) {
        this.port = port;
        this.serviceRegistry = serviceRegistry;
        this.serializer = serializer;
        this.threadPool = Executors.newFixedThreadPool(Runtime.getRuntime().availableProcessors() * 2);
        this.connectionStates = new ConcurrentHashMap<>();
    }

    @Override
    public void start() {
        try {
            selector = Selector.open();
            serverSocketChannel = ServerSocketChannel.open();
            serverSocketChannel.configureBlocking(false);
            serverSocketChannel.bind(new InetSocketAddress(port));
            serverSocketChannel.register(selector, SelectionKey.OP_ACCEPT);

            running = true;
            logger.info("Server started and listening on port {}", port);

            while (running) {
                selector.select();
                Iterator<SelectionKey> iter = selector.selectedKeys().iterator();

                while (iter.hasNext()) {
                    SelectionKey key = iter.next();
                    iter.remove();

                    try {
                        if (key.isAcceptable()) {
                            handleAccept(key);
                        } else if (key.isReadable()) {
                            handleRead(key);
                        }
                    } catch (Exception e) {
                        logger.error("Error handling key", e);
                        closeChannel(key);
                    }
                }
            }
        } catch (IOException e) {
            if (running) {
                logger.error("Server error", e);
                stop();
            }
        }
    }

    private void handleAccept(SelectionKey key) throws IOException {
        ServerSocketChannel ssc = (ServerSocketChannel) key.channel();
        SocketChannel clientChannel = ssc.accept();
        clientChannel.configureBlocking(false);

        // 注册读事件并关联连接状态
        ConnectionState state = new ConnectionState();
        connectionStates.put(clientChannel, state);
        clientChannel.register(selector, SelectionKey.OP_READ, state);

        logger.info("New client connected: {}", clientChannel.getRemoteAddress());
    }

    private void handleRead(SelectionKey key) throws IOException {
        SocketChannel clientChannel = (SocketChannel) key.channel();
        ConnectionState state = (ConnectionState) key.attachment();

        // 读取长度字段（如果尚未读取）
        if (state.lengthBuffer == null) {
            state.lengthBuffer = ByteBuffer.allocate(4);
        }

        if (clientChannel.read(state.lengthBuffer) == -1) {
            // 客户端关闭连接
            closeChannel(key);
            return;
        }

        if (state.lengthBuffer.hasRemaining()) {
            // 长度字段未读取完整
            return;
        }

        // 解析长度
        state.lengthBuffer.flip();
        int messageLength = state.lengthBuffer.getInt();
        state.lengthBuffer = null;

        // 读取消息体
        if (state.messageBuffer == null) {
            state.messageBuffer = ByteBuffer.allocate(messageLength);
        }

        int bytesRead = clientChannel.read(state.messageBuffer);
        if (bytesRead == -1) {
            closeChannel(key);
            return;
        }

        if (state.messageBuffer.hasRemaining()) {
            // 消息体未读取完整
            return;
        }

        // 消息读取完成，处理请求
        state.messageBuffer.flip();
        byte[] requestData = new byte[messageLength];
        state.messageBuffer.get(requestData);
        state.messageBuffer = null;

        // 提交到线程池处理
        long requestId = requestCounter.incrementAndGet();
        threadPool.submit(() -> processRequest(clientChannel, requestData, requestId));
    }

    private void processRequest(SocketChannel clientChannel, byte[] requestData, long requestId) {
        try {
            RpcRequest request = serializer.deserialize(requestData, RpcRequest.class);
            logger.info("Processing request {}: {}", requestId, request.getRequestId());

            RpcResponse response = new RpcResponse();
            response.setRequestId(request.getRequestId());

            try {
                Object result = handleRequest(request);
                response.setResult(result);
            } catch (Exception e) {
                logger.error("Error processing request {}", requestId, e);
                response.setError(e);
            }

            // 发送响应
            sendResponse(clientChannel, response, requestId);
        } catch (Exception e) {
            logger.error("Error deserializing request {}", requestId, e);
            try {
                clientChannel.close();
            } catch (IOException ex) {
                logger.error("Error closing channel", ex);
            }
        }
    }

    private void sendResponse(SocketChannel clientChannel, RpcResponse response, long requestId) {
        try {
            byte[] responseData = serializer.serialize(response);
            ByteBuffer buffer = ByteBuffer.allocate(4 + responseData.length);
            buffer.putInt(responseData.length);
            buffer.put(responseData);
            buffer.flip();

            while (buffer.hasRemaining()) {
                clientChannel.write(buffer);
            }

            logger.info("Response sent for request {}", requestId);
        } catch (IOException e) {
            logger.error("Failed to send response for request {}", requestId, e);
            try {
                clientChannel.close();
            } catch (IOException ex) {
                logger.error("Error closing channel", ex);
            }
        }
    }

    private void closeChannel(SelectionKey key) {
        try {
            SocketChannel channel = (SocketChannel) key.channel();
            logger.info("Closing channel: {}", channel.getRemoteAddress());
            connectionStates.remove(channel);
            key.cancel();
            channel.close();
        } catch (IOException e) {
            logger.error("Error closing channel", e);
        }
    }

    private Object handleRequest(RpcRequest request) throws NoSuchMethodException, InvocationTargetException, IllegalAccessException {
        // 获取服务实例
        Object service = serviceRegistry.getService(request.getServiceName());

        // 获取方法
        Method method = service.getClass().getMethod(
                request.getMethodName(),
                request.getParameterTypes());

        // 调用方法
        return method.invoke(service, request.getParameters());
    }

    @Override
    public void stop() {
        running = false;

        try {
            if (selector != null) {
                selector.close();
            }

            if (serverSocketChannel != null) {
                serverSocketChannel.close();
            }

            // 关闭所有客户端连接
            for (SocketChannel channel : connectionStates.keySet()) {
                try {
                    channel.close();
                } catch (IOException e) {
                    logger.error("Error closing client channel", e);
                }
            }

            connectionStates.clear();

            // 关闭线程池
            threadPool.shutdown();
            if (!threadPool.awaitTermination(5, TimeUnit.SECONDS)) {
                threadPool.shutdownNow();
            }

            logger.info("Server stopped");
        } catch (Exception e) {
            logger.error("Error stopping server", e);
        }
    }

    @Override
    public void registerService(Class<?> serviceInterface, Object serviceImpl) {
        serviceRegistry.registerService(serviceInterface, serviceImpl);
    }

    @Override
    public ServiceRegistry getServiceRegistry() {
        return serviceRegistry;
    }

    /**
     * 连接状态类
     */
    private static class ConnectionState {
        ByteBuffer lengthBuffer;
        ByteBuffer messageBuffer;
    }
}